{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "NQL Gridworld Pathfollowing",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "trEj8P0a7rET",
        "colab_type": "text"
      },
      "source": [
        "# NQL Gridworld Pathfollowing"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HdA4pjyL8BBm",
        "colab_type": "text"
      },
      "source": [
        "# Setup"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_gOkq73pft7S",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "outputId": "efe5e698-3b10-4685-9948-a0628588e200"
      },
      "source": [
        "%tensorflow_version 2.x\n",
        "!test -d language_repo || git clone https://github.com/google-research/language.git language_repo\n",
        "%cd /content/language_repo/language/nql\n",
        "!pip install ."
      ],
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Cloning into 'language_repo'...\n",
            "remote: Enumerating objects: 283, done.\u001b[K\n",
            "remote: Counting objects: 100% (283/283), done.\u001b[K\n",
            "remote: Compressing objects: 100% (221/221), done.\u001b[K\n",
            "remote: Total 1624 (delta 94), reused 175 (delta 59), pack-reused 1341\u001b[K\n",
            "Receiving objects: 100% (1624/1624), 1.99 MiB | 5.94 MiB/s, done.\n",
            "Resolving deltas: 100% (899/899), done.\n",
            "/content/language_repo/language/nql\n",
            "Processing /content/language_repo/language/nql\n",
            "Collecting tensorflow-gpu\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/0f/11/763f55d3d15efd778ef24453f126e6c33635680e5a2bb346da3fab5997cb/tensorflow_gpu-2.3.0-cp36-cp36m-manylinux2010_x86_64.whl (320.4MB)\n",
            "\u001b[K     |████████████████████████████████| 320.4MB 45kB/s \n",
            "\u001b[?25hRequirement already satisfied: scipy in /usr/local/lib/python3.6/dist-packages (from nql==0.0.1.dev0) (1.4.1)\n",
            "Collecting mock\n",
            "  Downloading https://files.pythonhosted.org/packages/cd/74/d72daf8dff5b6566db857cfd088907bb0355f5dd2914c4b3ef065c790735/mock-4.0.2-py3-none-any.whl\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.6/dist-packages (from nql==0.0.1.dev0) (1.18.5)\n",
            "Requirement already satisfied: astunparse==1.6.3 in /usr/local/lib/python3.6/dist-packages (from tensorflow-gpu->nql==0.0.1.dev0) (1.6.3)\n",
            "Requirement already satisfied: absl-py>=0.7.0 in /usr/local/lib/python3.6/dist-packages (from tensorflow-gpu->nql==0.0.1.dev0) (0.8.1)\n",
            "Requirement already satisfied: keras-preprocessing<1.2,>=1.1.1 in /usr/local/lib/python3.6/dist-packages (from tensorflow-gpu->nql==0.0.1.dev0) (1.1.2)\n",
            "Requirement already satisfied: opt-einsum>=2.3.2 in /usr/local/lib/python3.6/dist-packages (from tensorflow-gpu->nql==0.0.1.dev0) (3.3.0)\n",
            "Requirement already satisfied: wrapt>=1.11.1 in /usr/local/lib/python3.6/dist-packages (from tensorflow-gpu->nql==0.0.1.dev0) (1.12.1)\n",
            "Requirement already satisfied: wheel>=0.26 in /usr/local/lib/python3.6/dist-packages (from tensorflow-gpu->nql==0.0.1.dev0) (0.35.1)\n",
            "Requirement already satisfied: six>=1.12.0 in /usr/local/lib/python3.6/dist-packages (from tensorflow-gpu->nql==0.0.1.dev0) (1.15.0)\n",
            "Requirement already satisfied: grpcio>=1.8.6 in /usr/local/lib/python3.6/dist-packages (from tensorflow-gpu->nql==0.0.1.dev0) (1.31.0)\n",
            "Requirement already satisfied: protobuf>=3.9.2 in /usr/local/lib/python3.6/dist-packages (from tensorflow-gpu->nql==0.0.1.dev0) (3.12.4)\n",
            "Requirement already satisfied: h5py<2.11.0,>=2.10.0 in /usr/local/lib/python3.6/dist-packages (from tensorflow-gpu->nql==0.0.1.dev0) (2.10.0)\n",
            "Requirement already satisfied: tensorboard<3,>=2.3.0 in /usr/local/lib/python3.6/dist-packages (from tensorflow-gpu->nql==0.0.1.dev0) (2.3.0)\n",
            "Requirement already satisfied: termcolor>=1.1.0 in /usr/local/lib/python3.6/dist-packages (from tensorflow-gpu->nql==0.0.1.dev0) (1.1.0)\n",
            "Requirement already satisfied: tensorflow-estimator<2.4.0,>=2.3.0 in /usr/local/lib/python3.6/dist-packages (from tensorflow-gpu->nql==0.0.1.dev0) (2.3.0)\n",
            "Requirement already satisfied: gast==0.3.3 in /usr/local/lib/python3.6/dist-packages (from tensorflow-gpu->nql==0.0.1.dev0) (0.3.3)\n",
            "Requirement already satisfied: google-pasta>=0.1.8 in /usr/local/lib/python3.6/dist-packages (from tensorflow-gpu->nql==0.0.1.dev0) (0.2.0)\n",
            "Requirement already satisfied: setuptools in /usr/local/lib/python3.6/dist-packages (from protobuf>=3.9.2->tensorflow-gpu->nql==0.0.1.dev0) (49.6.0)\n",
            "Requirement already satisfied: google-auth<2,>=1.6.3 in /usr/local/lib/python3.6/dist-packages (from tensorboard<3,>=2.3.0->tensorflow-gpu->nql==0.0.1.dev0) (1.17.2)\n",
            "Requirement already satisfied: tensorboard-plugin-wit>=1.6.0 in /usr/local/lib/python3.6/dist-packages (from tensorboard<3,>=2.3.0->tensorflow-gpu->nql==0.0.1.dev0) (1.7.0)\n",
            "Requirement already satisfied: requests<3,>=2.21.0 in /usr/local/lib/python3.6/dist-packages (from tensorboard<3,>=2.3.0->tensorflow-gpu->nql==0.0.1.dev0) (2.23.0)\n",
            "Requirement already satisfied: google-auth-oauthlib<0.5,>=0.4.1 in /usr/local/lib/python3.6/dist-packages (from tensorboard<3,>=2.3.0->tensorflow-gpu->nql==0.0.1.dev0) (0.4.1)\n",
            "Requirement already satisfied: markdown>=2.6.8 in /usr/local/lib/python3.6/dist-packages (from tensorboard<3,>=2.3.0->tensorflow-gpu->nql==0.0.1.dev0) (3.2.2)\n",
            "Requirement already satisfied: werkzeug>=0.11.15 in /usr/local/lib/python3.6/dist-packages (from tensorboard<3,>=2.3.0->tensorflow-gpu->nql==0.0.1.dev0) (1.0.1)\n",
            "Requirement already satisfied: pyasn1-modules>=0.2.1 in /usr/local/lib/python3.6/dist-packages (from google-auth<2,>=1.6.3->tensorboard<3,>=2.3.0->tensorflow-gpu->nql==0.0.1.dev0) (0.2.8)\n",
            "Requirement already satisfied: rsa<5,>=3.1.4; python_version >= \"3\" in /usr/local/lib/python3.6/dist-packages (from google-auth<2,>=1.6.3->tensorboard<3,>=2.3.0->tensorflow-gpu->nql==0.0.1.dev0) (4.6)\n",
            "Requirement already satisfied: cachetools<5.0,>=2.0.0 in /usr/local/lib/python3.6/dist-packages (from google-auth<2,>=1.6.3->tensorboard<3,>=2.3.0->tensorflow-gpu->nql==0.0.1.dev0) (4.1.1)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.6/dist-packages (from requests<3,>=2.21.0->tensorboard<3,>=2.3.0->tensorflow-gpu->nql==0.0.1.dev0) (2020.6.20)\n",
            "Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /usr/local/lib/python3.6/dist-packages (from requests<3,>=2.21.0->tensorboard<3,>=2.3.0->tensorflow-gpu->nql==0.0.1.dev0) (1.24.3)\n",
            "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.6/dist-packages (from requests<3,>=2.21.0->tensorboard<3,>=2.3.0->tensorflow-gpu->nql==0.0.1.dev0) (2.10)\n",
            "Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.6/dist-packages (from requests<3,>=2.21.0->tensorboard<3,>=2.3.0->tensorflow-gpu->nql==0.0.1.dev0) (3.0.4)\n",
            "Requirement already satisfied: requests-oauthlib>=0.7.0 in /usr/local/lib/python3.6/dist-packages (from google-auth-oauthlib<0.5,>=0.4.1->tensorboard<3,>=2.3.0->tensorflow-gpu->nql==0.0.1.dev0) (1.3.0)\n",
            "Requirement already satisfied: importlib-metadata; python_version < \"3.8\" in /usr/local/lib/python3.6/dist-packages (from markdown>=2.6.8->tensorboard<3,>=2.3.0->tensorflow-gpu->nql==0.0.1.dev0) (1.7.0)\n",
            "Requirement already satisfied: pyasn1<0.5.0,>=0.4.6 in /usr/local/lib/python3.6/dist-packages (from pyasn1-modules>=0.2.1->google-auth<2,>=1.6.3->tensorboard<3,>=2.3.0->tensorflow-gpu->nql==0.0.1.dev0) (0.4.8)\n",
            "Requirement already satisfied: oauthlib>=3.0.0 in /usr/local/lib/python3.6/dist-packages (from requests-oauthlib>=0.7.0->google-auth-oauthlib<0.5,>=0.4.1->tensorboard<3,>=2.3.0->tensorflow-gpu->nql==0.0.1.dev0) (3.1.0)\n",
            "Requirement already satisfied: zipp>=0.5 in /usr/local/lib/python3.6/dist-packages (from importlib-metadata; python_version < \"3.8\"->markdown>=2.6.8->tensorboard<3,>=2.3.0->tensorflow-gpu->nql==0.0.1.dev0) (3.1.0)\n",
            "Building wheels for collected packages: nql\n",
            "  Building wheel for nql (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for nql: filename=nql-0.0.1.dev0-cp36-none-any.whl size=48738 sha256=ef554fd242a328c0bd5609ea788d22bd59987cb8c297d80fe759b6dfdd422b2d\n",
            "  Stored in directory: /tmp/pip-ephem-wheel-cache-ok3a8j9s/wheels/51/1e/3d/f92e27698ae7f6d586fd5482f569704ed8ba1a16f293bb0a5c\n",
            "Successfully built nql\n",
            "Installing collected packages: tensorflow-gpu, mock, nql\n",
            "Successfully installed mock-4.0.2 nql-0.0.1.dev0 tensorflow-gpu-2.3.0\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "d0l8oN-G7qVN",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "e873cc97-48b4-431c-9406-b59b0d4a5b60"
      },
      "source": [
        "import random\n",
        "import os\n",
        "\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "\n",
        "from tensorflow import keras\n",
        "from tensorflow.keras import layers, models, regularizers\n",
        "from tensorflow.keras.optimizers import schedules\n",
        "from tensorflow.keras.preprocessing import text\n",
        "from tensorflow.keras.models import Sequential\n",
        "\n",
        "import nql\n",
        "\n",
        "print(tf.__version__)"
      ],
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "2.3.0\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Q65_-hBx3t41",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Makes things as deterministic as possible for reproducibility\n",
        "seed = 1234 #@param{type: \"integer\"}\n",
        "tf.random.set_seed(seed)\n",
        "random.seed(seed)"
      ],
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5Hk02Z5yBuqM",
        "colab_type": "text"
      },
      "source": [
        "## Build the gridworld"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "48BJzTQWF3hr",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class GridWorld():\n",
        "  \"\"\"A simple grid world for testing NQL algorithms.\n",
        "  \n",
        "  The world is square grid with every cell connected to its cardinal direction\n",
        "  neighbors, with the exception of several holes which are inescapable.\"\"\"\n",
        "  size = None\n",
        "  context = None\n",
        "  words = ['go', 'top', 'left', 'right', 'bottom', 'center', 'up', 'down',\n",
        "           'then']\n",
        "  # A 2D dictionary mapping a cell and a valid move to another cell.\n",
        "  # Ex: _cell_move_cell['cell_2_3']['up'] == 'cell_1_3'\n",
        "  _cell_move_cell = {}\n",
        "  # _cell_paths[start_cell][num_moves][end_cell] is a list of efficient paths\n",
        "  # starting at start_cell, taking num_moves, and ending at end_cell\n",
        "  _cell_paths = {}\n",
        "  # _cell_dists[start_cell][dist] is a set of all of the cells exactly dist away\n",
        "  # from start cell\n",
        "  _cell_dists = {}\n",
        "  \n",
        "  def __init__(self, size, n_holes):\n",
        "    if size ** 2 < n_holes:\n",
        "      raise ValueError(\n",
        "          'size ^ 2 (%d) must be great than n_holes (%d).'\n",
        "          % (size ** 2, n_holes)) \n",
        "    self.size = size\n",
        "    self.holes = set()\n",
        "    \n",
        "    # Remove cells\n",
        "    while len(self.holes) < n_holes:\n",
        "      r = random.randrange(size)\n",
        "      c = random.randrange(size)\n",
        "      self.holes.add((r,c))\n",
        "    \n",
        "    self.context = nql.NeuralQueryContext()\n",
        "    self.context.declare_relation('up','place_t','place_t')\n",
        "    self.context.declare_relation('down','place_t','place_t')\n",
        "    self.context.declare_relation('left','place_t','place_t')\n",
        "    self.context.declare_relation('right','place_t','place_t')\n",
        "    kg_lines = []\n",
        "    for (r,c) in [(r,c) for r in range(self.size) for c in range(self.size)]:\n",
        "      if (r,c) in self.holes:\n",
        "        continue\n",
        "      self._cell_move_cell[self._cell(r,c)] = {}\n",
        "      def connect_to_cell(dest_r, dest_c, dir):\n",
        "        dest = self._cell(dest_r, dest_c)\n",
        "        self._cell_move_cell[self._cell(r,c)][dir] = dest\n",
        "        kg_lines.append('%s\\t%s\\t%s' % (dir, self._cell(r,c), dest))\n",
        "      if(r > 0):\n",
        "        connect_to_cell(r-1, c, 'up')\n",
        "      if(r < self.size-1):\n",
        "        connect_to_cell(r+1, c, 'down')\n",
        "      if(c > 0):\n",
        "        connect_to_cell(r, c-1, 'left')\n",
        "      if(c < self.size-1):\n",
        "        connect_to_cell(r, c+1, 'right')\n",
        "\n",
        "    self.context.load_kg(lines=kg_lines, freeze=True)\n",
        "    self.context.construct_relation_group('dir_g', 'place_t', 'place_t')\n",
        "\n",
        "  def _cell(self, i, j):\n",
        "    if (i, j) in self.holes:\n",
        "      return 'hole_%d_%d' % (i+1, j+1)\n",
        "    return 'cell_%d_%d' % (i+1, j+1)\n",
        "\n",
        "  # Output: [(query, starting_cell, ending_cell, num moves)]\n",
        "  def generate_examples(self, num, possible_moves=[1,2,3]):\n",
        "    # Extends all paths from start_cell by 1 step\n",
        "    def _extend_paths(start_cell):\n",
        "      cell_paths = self._cell_paths[start_cell]\n",
        "      cell_dists = self._cell_dists[start_cell]\n",
        "      prev_path_len = len(cell_dists) - 1\n",
        "\n",
        "      cell_dists.append(set())\n",
        "      cell_paths.append({})\n",
        "      if not len(cell_dists[prev_path_len]):\n",
        "        # We're already longer than the longest path possible from this cell.\n",
        "        return\n",
        "  \n",
        "      seen_cells = set(\n",
        "          [cell for cells_at_dist in cell_dists for cell in cells_at_dist])\n",
        "      for cell in cell_dists[prev_path_len]:\n",
        "        prev_paths = cell_paths[prev_path_len][cell]\n",
        "        if cell not in self._cell_move_cell:\n",
        "          continue # This is a hole\n",
        "        for (dir, next_cell) in self._cell_move_cell[cell].items():\n",
        "          if next_cell not in seen_cells: # This is an efficient path\n",
        "            if next_cell not in cell_dists[prev_path_len+1]:\n",
        "              # This is the first time we found this cell\n",
        "              cell_dists[prev_path_len+1].add(next_cell)\n",
        "              cell_paths[prev_path_len+1][next_cell] = []\n",
        "            for path in prev_paths:\n",
        "              if prev_path_len == 0:\n",
        "                new_path = 'go ' + dir\n",
        "              else:\n",
        "                new_path = path + ' then ' + dir\n",
        "              cell_paths[prev_path_len+1][next_cell].append(new_path)\n",
        "    def generate_example(num_moves):\n",
        "      example = 'go '\n",
        "      while True:\n",
        "        start_row = random.randrange(self.size)\n",
        "        start_col = random.randrange(self.size)\n",
        "        starting_cell = self._cell(start_row, start_col)\n",
        "        if not starting_cell in self._cell_move_cell:\n",
        "          # Starting cell is a hole. Try again.\n",
        "          continue\n",
        "        if not starting_cell in self._cell_paths:\n",
        "          # We've never started from this cell before.\n",
        "          self._cell_paths[starting_cell] = [{starting_cell: ['']}]\n",
        "          self._cell_dists[starting_cell] = [{starting_cell}]\n",
        "        while len(self._cell_paths[starting_cell]) <= num_moves:\n",
        "          _extend_paths(starting_cell)\n",
        "        possible_ending_cells = list(self._cell_paths[starting_cell][num_moves])\n",
        "        if not possible_ending_cells:\n",
        "          # No paths num_moves long from this cell. Start over.\n",
        "          continue\n",
        "        ending_cell = random.choice(possible_ending_cells)\n",
        "        path = random.choice(\n",
        "            self._cell_paths[starting_cell][num_moves][ending_cell])\n",
        "        return (path, starting_cell, ending_cell)\n",
        "\n",
        "    examples = [None] * num\n",
        "    while True:\n",
        "      for i in range(num):\n",
        "        # Each element of possible_moves is equally likely to be selected\n",
        "        num_moves = random.choice(possible_moves)\n",
        "        example = ''\n",
        "        while not example:\n",
        "          example, starting_cell, ending_cell = generate_example(num_moves)\n",
        "        if examples[i] is None:\n",
        "          examples[i] = ({\n",
        "            \"input_text\": example,\n",
        "            \"start\": tf.squeeze(self.context.one(starting_cell, 'place_t').tf),\n",
        "            \"start_name\": starting_cell,\n",
        "            \"end_name\": ending_cell,\n",
        "            }, tf.squeeze(self.context.one(ending_cell, 'place_t').tf))\n",
        "        yield examples[i]"
      ],
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NvR0gt7KBtjR",
        "colab_type": "code",
        "cellView": "both",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 54
        },
        "outputId": "5f1f0849-ba67-4d43-b364-cb5df72828b4"
      },
      "source": [
        "grid_size =  10#@param {type: \"integer\"}\n",
        "p_holes = 0.25 #@param {type: \"number\"}\n",
        "grid_holes = grid_size ** 2 * p_holes\n",
        "max_moves =  10#@param {type: \"integer\"}\n",
        "\n",
        "env = GridWorld(grid_size, grid_holes)\n",
        "print(\"Holes:\", env.holes)"
      ],
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Holes: {(7, 3), (4, 7), (9, 1), (2, 1), (8, 9), (9, 0), (3, 3), (8, 1), (7, 6), (0, 4), (1, 1), (9, 7), (5, 4), (0, 0), (7, 1), (5, 2), (0, 5), (1, 9), (1, 0), (0, 8), (5, 3), (0, 1), (2, 0), (1, 8), (7, 8)}\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fq1lh0Vl0VVL",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 748
        },
        "outputId": "f37d6e21-3957-4c8e-fd8d-4ecf4b7fa333"
      },
      "source": [
        "from itertools import islice\n",
        "print(\"Examples:\", list(islice(env.generate_examples(4), 3)))\n",
        "nqc = env.context"
      ],
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Examples: [({'input_text': 'go down then right', 'start': <tf.Tensor: shape=(99,), dtype=float32, numpy=\n",
            "array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
            "       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
            "       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
            "       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
            "       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
            "       0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
            "      dtype=float32)>, 'start_name': 'cell_9_4', 'end_name': 'cell_10_5'}, <tf.Tensor: shape=(99,), dtype=float32, numpy=\n",
            "array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
            "       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
            "       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
            "       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
            "       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
            "       0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
            "      dtype=float32)>), ({'input_text': 'go up then up', 'start': <tf.Tensor: shape=(99,), dtype=float32, numpy=\n",
            "array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
            "       0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
            "       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
            "       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
            "       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
            "       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
            "      dtype=float32)>, 'start_name': 'cell_3_5', 'end_name': 'hole_1_5'}, <tf.Tensor: shape=(99,), dtype=float32, numpy=\n",
            "array([0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
            "       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
            "       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
            "       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
            "       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
            "       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
            "      dtype=float32)>), ({'input_text': 'go down then right', 'start': <tf.Tensor: shape=(99,), dtype=float32, numpy=\n",
            "array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
            "       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
            "       0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
            "       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
            "       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
            "       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
            "      dtype=float32)>, 'start_name': 'cell_5_1', 'end_name': 'cell_6_2'}, <tf.Tensor: shape=(99,), dtype=float32, numpy=\n",
            "array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
            "       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
            "       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0.,\n",
            "       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
            "       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
            "       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
            "      dtype=float32)>)]\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "960yzsxf-jsP",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "9bd3991c-3e53-41fc-984f-797fa3b74e8a"
      },
      "source": [
        "dataset_size = 100000 #@param{type: \"integer\"}\n",
        "train_gen = env.generate_examples(num=dataset_size, possible_moves=range(1, max_moves+1))\n",
        "test_gen = env.generate_examples(int(dataset_size / 10))\n",
        "max_seq_len = max(len(x[\"input_text\"].split(\" \")) for (x, _) in islice(train_gen, 1000))\n",
        "print(\"max_seq_len is %d\" % max_seq_len)\n",
        "assert max_seq_len >= 2\n",
        "assert max_seq_len < 50"
      ],
      "execution_count": 7,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "max_seq_len is 20\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OtO0SZbckhWR",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "train_dataset = tf.data.Dataset.from_generator(\n",
        "    lambda: (x for x in train_gen),\n",
        "    output_types=(\n",
        "        { \"input_text\": tf.string,\n",
        "          \"start\": tf.float32,\n",
        "          \"start_name\": tf.string,\n",
        "          \"end_name\": tf.string,\n",
        "        }, tf.float32),\n",
        "    output_shapes=(\n",
        "        { \"input_text\": tf.TensorShape([]),\n",
        "          \"start\": tf.TensorShape([nqc.get_max_id('place_t')]),\n",
        "          \"start_name\": tf.TensorShape([]),\n",
        "          \"end_name\": tf.TensorShape([]),\n",
        "        }, tf.TensorShape([nqc.get_max_id('place_t')])),\n",
        "    )\n",
        "text_dataset = train_dataset.map(lambda x,_: x[\"input_text\"])"
      ],
      "execution_count": 8,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "i5NO4-5i2Lir",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 649
        },
        "outputId": "febc0cc3-4219-4a93-ce95-4b42ca5259de"
      },
      "source": [
        "for i in train_dataset.take(2):\n",
        "  print(repr(i))\n",
        "for i in text_dataset.take(2):\n",
        "  print(repr(i))\n",
        "print(repr(train_dataset))\n",
        "print(repr(train_dataset.element_spec))\n",
        "print(repr(text_dataset))\n",
        "print(repr(text_dataset.element_spec))"
      ],
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "({'input_text': <tf.Tensor: shape=(), dtype=string, numpy=b'go down then down then down then right then down then right'>, 'start': <tf.Tensor: shape=(99,), dtype=float32, numpy=\n",
            "array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
            "       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
            "       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
            "       0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
            "       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
            "       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
            "      dtype=float32)>, 'start_name': <tf.Tensor: shape=(), dtype=string, numpy=b'cell_6_6'>, 'end_name': <tf.Tensor: shape=(), dtype=string, numpy=b'hole_10_8'>}, <tf.Tensor: shape=(99,), dtype=float32, numpy=\n",
            "array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
            "       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
            "       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
            "       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
            "       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
            "       0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0.],\n",
            "      dtype=float32)>)\n",
            "({'input_text': <tf.Tensor: shape=(), dtype=string, numpy=b'go down then right'>, 'start': <tf.Tensor: shape=(99,), dtype=float32, numpy=\n",
            "array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
            "       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
            "       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
            "       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
            "       0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
            "       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
            "      dtype=float32)>, 'start_name': <tf.Tensor: shape=(), dtype=string, numpy=b'cell_8_5'>, 'end_name': <tf.Tensor: shape=(), dtype=string, numpy=b'cell_9_6'>}, <tf.Tensor: shape=(99,), dtype=float32, numpy=\n",
            "array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
            "       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
            "       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
            "       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
            "       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0.,\n",
            "       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
            "      dtype=float32)>)\n",
            "<tf.Tensor: shape=(), dtype=string, numpy=b'go left'>\n",
            "<tf.Tensor: shape=(), dtype=string, numpy=b'go right'>\n",
            "<FlatMapDataset shapes: ({input_text: (), start: (99,), start_name: (), end_name: ()}, (99,)), types: ({input_text: tf.string, start: tf.float32, start_name: tf.string, end_name: tf.string}, tf.float32)>\n",
            "({'input_text': TensorSpec(shape=(), dtype=tf.string, name=None), 'start': TensorSpec(shape=(99,), dtype=tf.float32, name=None), 'start_name': TensorSpec(shape=(), dtype=tf.string, name=None), 'end_name': TensorSpec(shape=(), dtype=tf.string, name=None)}, TensorSpec(shape=(99,), dtype=tf.float32, name=None))\n",
            "<MapDataset shapes: (), types: tf.string>\n",
            "TensorSpec(shape=(), dtype=tf.string, name=None)\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XmZoWsWDE8UF",
        "colab_type": "text"
      },
      "source": [
        "## Example"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "_1hWBl2ti_U6",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 85
        },
        "outputId": "f95fc6db-f436-4a6e-9243-bc3b8cae5b8c"
      },
      "source": [
        "name = nqc.get_entity_name(grid_size * 2, 'place_t')\n",
        "print(name)\n",
        "cell = nqc.one(name, 'place_t')\n",
        "print(cell.eval())\n",
        "cell = nqc.one(name, 'place_t').follow('right')\n",
        "print(cell.eval())\n",
        "cell = nqc.one(name, 'place_t').follow('right').follow('up')\n",
        "print(cell.eval())"
      ],
      "execution_count": 10,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "cell_3_6\n",
            "{'cell_3_6': 1.0}\n",
            "{'cell_3_7': 1.0}\n",
            "{'cell_2_7': 1.0}\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0UTFDiJ96Rb4",
        "colab_type": "text"
      },
      "source": [
        "# Multiple Moves: Solve a problem with multiple possible templates"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vzlpJ-jD2n-F",
        "colab_type": "text"
      },
      "source": [
        "## RNN Model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lzzRS554lzyf",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "#from tensorflow.compat.v2.keras.layers.experimental.preprocessing import TextVectorization\n",
        "from tensorflow.keras.layers.experimental.preprocessing import TextVectorization\n",
        "\n",
        "class NqlFollowLayer(layers.Layer):\n",
        "  def __init__(self, context, **kwargs):\n",
        "    self.context = context\n",
        "    super(NqlFollowLayer, self).__init__(**kwargs)\n",
        "  \n",
        "  def build(self, input_shape):\n",
        "    assert isinstance(input_shape, list)\n",
        "    super(NqlFollowLayer, self).build(input_shape)\n",
        "  \n",
        "  def call(self, x):\n",
        "    assert isinstance(x, list)\n",
        "    place_tf, dir_tf = x\n",
        "    place_nql = self.context.as_nql(place_tf, \"place_t\")\n",
        "    assert isinstance(place_nql, nql.NeuralQueryExpression)\n",
        "    dir_nql = self.context.as_nql(dir_tf, \"dir_g\")\n",
        "    assert isinstance(dir_nql, nql.NeuralQueryExpression)\n",
        "    new_place_nql = place_nql.follow(dir_nql)\n",
        "    return new_place_nql.tf\n",
        "  \n",
        "  def compute_output_shape(self, input_shape):\n",
        "    return self.context.get_max_id(\"place_t\")\n",
        "  \n",
        "  def get_config(self):\n",
        "    config = super(NqlFollowLayer, self).get_config()\n",
        "    config['context'] = self.context\n",
        "    return config\n",
        "\n",
        "def build_model(nqc, layer_width, embedding_dim, max_num_moves, text_dataset, num_layers=1, dropout=0.0, l1=0.0001, l2=0.0001):\n",
        "  # Transforms text input into int sequences of length max_seq_len\n",
        "  vectorize_layer = TextVectorization(max_tokens=len(env.words),\n",
        "                                      output_mode=\"int\",\n",
        "                                      output_sequence_length=max_seq_len,\n",
        "                                      name=\"VectorizationLayer\",\n",
        "                                      )\n",
        "  vectorize_layer.adapt(text_dataset)\n",
        "\n",
        "  # Build the encoder\n",
        "  text_input = layers.Input(shape=(1,), dtype=tf.string, name=\"input_text\")\n",
        "  vec_input = vectorize_layer(text_input)\n",
        "  # Account for the OOV token in this layer\n",
        "  embedding_layer = layers.Embedding(input_dim=len(env.words)+1,\n",
        "                                     output_dim=embedding_dim,\n",
        "                                     mask_zero=True,\n",
        "                                     embeddings_regularizer=regularizers.l1_l2(l1=l1, l2=l2),\n",
        "                                     name=\"TextEmbedding\")\n",
        "  enc_emb = embedding_layer(vec_input)\n",
        "  # Use an LSTM layer here to process the whole input sequence\n",
        "  encoder_lstm = layers.LSTM(layer_width,\n",
        "                             recurrent_dropout=dropout,\n",
        "                             kernel_regularizer=regularizers.l1_l2(l1=l1, l2=l2),\n",
        "                             recurrent_regularizer=regularizers.l1_l2(l1=l1, l2=l2),\n",
        "                             bias_regularizer=regularizers.l1_l2(l1=l1, l2=l2),\n",
        "                             return_state=True,\n",
        "                             name=\"EncoderLstm\"\n",
        "                             )\n",
        "  enc_out, enc_state_h, enc_state_c = encoder_lstm(enc_emb)\n",
        "\n",
        "  # Build the decoder\n",
        "  decoder_lstm = layers.LSTM(layer_width,\n",
        "                             recurrent_dropout=dropout,\n",
        "                             kernel_regularizer=regularizers.l1_l2(l1=l1, l2=l2),\n",
        "                             recurrent_regularizer=regularizers.l1_l2(l1=l1, l2=l2),\n",
        "                             bias_regularizer=regularizers.l1_l2(l1=l1, l2=l2),\n",
        "                             return_sequences=True,\n",
        "                             name=\"DecoderLstm\"\n",
        "                             )\n",
        "\n",
        "  decoder_move_model = Sequential(name=\"DecoderMoveModel\")\n",
        "  for i in range(num_layers-1):\n",
        "    decoder_move_model.add(layers.Dense(layer_width,\n",
        "                                        kernel_regularizer=regularizers.l1_l2(l1=l1, l2=l2),\n",
        "                                        bias_regularizer=regularizers.l1_l2(l1=l1, l2=l2),\n",
        "                                        ))\n",
        "  decoder_move_model.add(layers.Dense(\n",
        "      nqc.get_max_id(\"dir_g\"), activation=\"softmax\"\n",
        "  ))\n",
        "\n",
        "  decoder_prob_model = Sequential(name=\"DecoderProbModel\")\n",
        "  for i in range(num_layers-1):\n",
        "    decoder_prob_model.add(layers.Dense(layer_width,\n",
        "                                        kernel_regularizer=regularizers.l1_l2(l1=l1, l2=l2),\n",
        "                                        bias_regularizer=regularizers.l1_l2(l1=l1, l2=l2),\n",
        "                                        ))\n",
        "  decoder_prob_model.add(layers.Dense(1, activation=\"sigmoid\"))\n",
        "\n",
        "  initial_place = layers.Input(shape=(nqc.get_max_id(\"place_t\"),), name=\"start\")\n",
        "  # Initial state of the decoder is the final state of the encoder\n",
        "  dec_states = [enc_state_h, enc_state_c]\n",
        "  cur_place = initial_place\n",
        "  nql_layer = NqlFollowLayer(nqc, dynamic=True)\n",
        "  # This should get turned into a tensor of the right shape the first time it's\n",
        "  # updated.\n",
        "  prob_remaining = 1.0\n",
        "  final_place = tf.zeros(shape=(nqc.get_max_id('place_t'),))\n",
        "  decoder_lstm_outs = decoder_lstm(\n",
        "      tf.ones(shape=(1, max_num_moves, layer_width)),\n",
        "      initial_state=dec_states)\n",
        "  for i in range(max_num_moves):\n",
        "    move_out = decoder_move_model(decoder_lstm_outs[:,i,:])\n",
        "    prob_out = decoder_prob_model(decoder_lstm_outs[:,i,:])\n",
        "\n",
        "    ## Use the output move to update the current place\n",
        "    cur_place = nql_layer([cur_place, move_out])\n",
        "    prob_stopping = prob_remaining * prob_out\n",
        "    prob_remaining = prob_remaining - prob_stopping\n",
        "\n",
        "    ## Update final output place\n",
        "    final_place = final_place + (prob_stopping * cur_place)\n",
        "  \n",
        "  # Build the final model that goes from text to a place\n",
        "  model = models.Model(inputs=[text_input, initial_place], outputs=final_place)\n",
        "  return model"
      ],
      "execution_count": 11,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ySKk2wHm2qQP",
        "colab_type": "code",
        "cellView": "both",
        "colab": {}
      },
      "source": [
        "#@title Model Params { form-width: \"25%\" }\n",
        "layer_width =  128 #@param{type: \"integer\"}\n",
        "embedding_dim =   4#@param{type: \"integer\"}\n",
        "num_layers =   3#@param{type: \"integer\"}\n",
        "l1 = 0.00002 #@param{type: \"number\"}\n",
        "l2 = 0.00002 #@param{type: \"number\"}\n",
        "dropout = 0.1 #@param{type: \"number\"}\n",
        "\n",
        "# For setting up the TextVectorization layer\n",
        "text_dataset_sample = text_dataset.take(1024).batch(1024)\n",
        "model = build_model(nqc,\n",
        "                    layer_width=layer_width,\n",
        "                    embedding_dim=embedding_dim,\n",
        "                    max_num_moves=max_moves,\n",
        "                    text_dataset=text_dataset_sample,\n",
        "                    dropout=dropout,\n",
        "                    l1=l1,\n",
        "                    l2=l2,\n",
        "                    num_layers=num_layers)"
      ],
      "execution_count": 12,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ijpLq-xZg5BC",
        "colab_type": "code",
        "cellView": "both",
        "colab": {}
      },
      "source": [
        "batch_size =  1024#@param{type: \"integer\"}\n",
        "learning_rate = 0.0025 #@param{type: \"number\"}\n",
        "clip_norm = 2.0 #@param{type: \"number\"}\n",
        "decay = 0.005 #@param{type: \"number\"}\n",
        "\n",
        "optimizer = tf.keras.optimizers.Nadam(learning_rate=learning_rate,\n",
        "                                      clipnorm=clip_norm,\n",
        "                                      decay=decay,\n",
        "                                     )\n",
        "\n",
        "model.compile(optimizer=optimizer,\n",
        "              loss=\"categorical_crossentropy\",\n",
        "              metrics=[\"accuracy\"],\n",
        "              # run_eagerly=True, # TODO(rofer): Figure out if this is only needed on CPUs\n",
        "              )\n",
        "batched_dataset = train_dataset.batch(batch_size)\n",
        "\n",
        "callbacks = [\n",
        "             keras.callbacks.TerminateOnNaN(),\n",
        "]"
      ],
      "execution_count": 13,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "g09QBCjtYE5K",
        "colab_type": "code",
        "cellView": "both",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "outputId": "0aaa85eb-2aba-40ca-939f-4b07e2a30572"
      },
      "source": [
        "n_epochs =  100#@param{type: \"integer\"}\n",
        "history = model.fit(batched_dataset,\n",
        "                    epochs=n_epochs,\n",
        "                    steps_per_epoch=50,\n",
        "                    callbacks=callbacks)"
      ],
      "execution_count": 14,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Epoch 1/100\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/functional.py:543: UserWarning: Input dict contained keys ['start_name', 'end_name'] which did not match any model input. They will be ignored by the model.\n",
            "  [n for n in tensors.keys() if n not in ref_input_names])\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "50/50 [==============================] - 103s 2s/step - loss: 6.1354 - accuracy: 0.0300\n",
            "Epoch 2/100\n",
            "50/50 [==============================] - 101s 2s/step - loss: 4.6080 - accuracy: 0.0939\n",
            "Epoch 3/100\n",
            "50/50 [==============================] - 89s 2s/step - loss: 2.7759 - accuracy: 0.2668\n",
            "Epoch 4/100\n",
            "50/50 [==============================] - 89s 2s/step - loss: 2.1506 - accuracy: 0.3858\n",
            "Epoch 5/100\n",
            "50/50 [==============================] - 89s 2s/step - loss: 1.8828 - accuracy: 0.4398\n",
            "Epoch 6/100\n",
            "50/50 [==============================] - 90s 2s/step - loss: 1.6598 - accuracy: 0.4837\n",
            "Epoch 7/100\n",
            "50/50 [==============================] - 94s 2s/step - loss: 1.4606 - accuracy: 0.5353\n",
            "Epoch 8/100\n",
            "50/50 [==============================] - 92s 2s/step - loss: 1.3069 - accuracy: 0.5760\n",
            "Epoch 9/100\n",
            "50/50 [==============================] - 92s 2s/step - loss: 1.1738 - accuracy: 0.6135\n",
            "Epoch 10/100\n",
            "50/50 [==============================] - 93s 2s/step - loss: 1.0895 - accuracy: 0.6317\n",
            "Epoch 11/100\n",
            "50/50 [==============================] - 91s 2s/step - loss: 0.9568 - accuracy: 0.6737\n",
            "Epoch 12/100\n",
            "50/50 [==============================] - 91s 2s/step - loss: 0.9902 - accuracy: 0.6731\n",
            "Epoch 13/100\n",
            "50/50 [==============================] - 94s 2s/step - loss: 0.8843 - accuracy: 0.7028\n",
            "Epoch 14/100\n",
            "50/50 [==============================] - 90s 2s/step - loss: 0.8141 - accuracy: 0.7344\n",
            "Epoch 15/100\n",
            "50/50 [==============================] - 88s 2s/step - loss: 0.8014 - accuracy: 0.7390\n",
            "Epoch 16/100\n",
            "50/50 [==============================] - 89s 2s/step - loss: 0.7513 - accuracy: 0.7597\n",
            "Epoch 17/100\n",
            "50/50 [==============================] - 91s 2s/step - loss: 0.7397 - accuracy: 0.7630\n",
            "Epoch 18/100\n",
            "50/50 [==============================] - 90s 2s/step - loss: 0.8688 - accuracy: 0.7356\n",
            "Epoch 19/100\n",
            "50/50 [==============================] - 89s 2s/step - loss: 0.6849 - accuracy: 0.7882\n",
            "Epoch 20/100\n",
            "50/50 [==============================] - 90s 2s/step - loss: 0.6643 - accuracy: 0.7956\n",
            "Epoch 21/100\n",
            "50/50 [==============================] - 88s 2s/step - loss: 0.6562 - accuracy: 0.7982\n",
            "Epoch 22/100\n",
            "50/50 [==============================] - 90s 2s/step - loss: 0.6091 - accuracy: 0.8192\n",
            "Epoch 23/100\n",
            "50/50 [==============================] - 91s 2s/step - loss: 0.6103 - accuracy: 0.8160\n",
            "Epoch 24/100\n",
            "50/50 [==============================] - 89s 2s/step - loss: 0.5798 - accuracy: 0.8313\n",
            "Epoch 25/100\n",
            "50/50 [==============================] - 89s 2s/step - loss: 0.5720 - accuracy: 0.8337\n",
            "Epoch 26/100\n",
            "50/50 [==============================] - 90s 2s/step - loss: 0.5694 - accuracy: 0.8323\n",
            "Epoch 27/100\n",
            "50/50 [==============================] - 89s 2s/step - loss: 0.5384 - accuracy: 0.8468\n",
            "Epoch 28/100\n",
            "50/50 [==============================] - 90s 2s/step - loss: 0.5290 - accuracy: 0.8468\n",
            "Epoch 29/100\n",
            "50/50 [==============================] - 89s 2s/step - loss: 0.5159 - accuracy: 0.8529\n",
            "Epoch 30/100\n",
            "50/50 [==============================] - 91s 2s/step - loss: 0.8437 - accuracy: 0.7740\n",
            "Epoch 31/100\n",
            "50/50 [==============================] - 89s 2s/step - loss: 0.4952 - accuracy: 0.8642\n",
            "Epoch 32/100\n",
            "50/50 [==============================] - 88s 2s/step - loss: 0.5155 - accuracy: 0.8543\n",
            "Epoch 33/100\n",
            "50/50 [==============================] - 89s 2s/step - loss: 0.4854 - accuracy: 0.8651\n",
            "Epoch 34/100\n",
            "50/50 [==============================] - 89s 2s/step - loss: 0.4777 - accuracy: 0.8634\n",
            "Epoch 35/100\n",
            "50/50 [==============================] - 88s 2s/step - loss: 0.4532 - accuracy: 0.8769\n",
            "Epoch 36/100\n",
            "50/50 [==============================] - 87s 2s/step - loss: 0.4660 - accuracy: 0.8698\n",
            "Epoch 37/100\n",
            "50/50 [==============================] - 89s 2s/step - loss: 0.4494 - accuracy: 0.8770\n",
            "Epoch 38/100\n",
            "50/50 [==============================] - 91s 2s/step - loss: 0.4534 - accuracy: 0.8755\n",
            "Epoch 39/100\n",
            "50/50 [==============================] - 92s 2s/step - loss: 0.4446 - accuracy: 0.8755\n",
            "Epoch 40/100\n",
            "50/50 [==============================] - 93s 2s/step - loss: 0.4490 - accuracy: 0.8742\n",
            "Epoch 41/100\n",
            "50/50 [==============================] - 89s 2s/step - loss: 0.4095 - accuracy: 0.8923\n",
            "Epoch 42/100\n",
            "50/50 [==============================] - 90s 2s/step - loss: 0.4285 - accuracy: 0.8808\n",
            "Epoch 43/100\n",
            "50/50 [==============================] - 90s 2s/step - loss: 0.4061 - accuracy: 0.8890\n",
            "Epoch 44/100\n",
            "50/50 [==============================] - 91s 2s/step - loss: 0.4263 - accuracy: 0.8787\n",
            "Epoch 45/100\n",
            "50/50 [==============================] - 90s 2s/step - loss: 0.3940 - accuracy: 0.8949\n",
            "Epoch 46/100\n",
            "50/50 [==============================] - 90s 2s/step - loss: 0.4208 - accuracy: 0.8788\n",
            "Epoch 47/100\n",
            "50/50 [==============================] - 89s 2s/step - loss: 0.3925 - accuracy: 0.8961\n",
            "Epoch 48/100\n",
            "50/50 [==============================] - 91s 2s/step - loss: 0.6693 - accuracy: 0.8373\n",
            "Epoch 49/100\n",
            "50/50 [==============================] - 93s 2s/step - loss: 0.5107 - accuracy: 0.8527\n",
            "Epoch 50/100\n",
            "50/50 [==============================] - 93s 2s/step - loss: 0.4022 - accuracy: 0.8928\n",
            "Epoch 51/100\n",
            "50/50 [==============================] - 92s 2s/step - loss: 0.3703 - accuracy: 0.9043\n",
            "Epoch 52/100\n",
            "50/50 [==============================] - 92s 2s/step - loss: 0.3998 - accuracy: 0.8883\n",
            "Epoch 53/100\n",
            "50/50 [==============================] - 94s 2s/step - loss: 0.3645 - accuracy: 0.9054\n",
            "Epoch 54/100\n",
            "50/50 [==============================] - 94s 2s/step - loss: 0.3891 - accuracy: 0.8941\n",
            "Epoch 55/100\n",
            "50/50 [==============================] - 92s 2s/step - loss: 0.3628 - accuracy: 0.9054\n",
            "Epoch 56/100\n",
            "50/50 [==============================] - 92s 2s/step - loss: 0.3748 - accuracy: 0.8970\n",
            "Epoch 57/100\n",
            "50/50 [==============================] - 93s 2s/step - loss: 0.3601 - accuracy: 0.9044\n",
            "Epoch 58/100\n",
            "50/50 [==============================] - 94s 2s/step - loss: 0.3684 - accuracy: 0.9002\n",
            "Epoch 59/100\n",
            "50/50 [==============================] - 97s 2s/step - loss: 0.3516 - accuracy: 0.9085\n",
            "Epoch 60/100\n",
            "50/50 [==============================] - 90s 2s/step - loss: 0.3655 - accuracy: 0.9015\n",
            "Epoch 61/100\n",
            "50/50 [==============================] - 90s 2s/step - loss: 0.3561 - accuracy: 0.9040\n",
            "Epoch 62/100\n",
            "50/50 [==============================] - 90s 2s/step - loss: 0.3539 - accuracy: 0.9054\n",
            "Epoch 63/100\n",
            "50/50 [==============================] - 91s 2s/step - loss: 0.3452 - accuracy: 0.9083\n",
            "Epoch 64/100\n",
            "50/50 [==============================] - 93s 2s/step - loss: 0.3389 - accuracy: 0.9113\n",
            "Epoch 65/100\n",
            "50/50 [==============================] - 92s 2s/step - loss: 0.3303 - accuracy: 0.9135\n",
            "Epoch 66/100\n",
            "50/50 [==============================] - 93s 2s/step - loss: 0.3541 - accuracy: 0.9051\n",
            "Epoch 67/100\n",
            "50/50 [==============================] - 91s 2s/step - loss: 0.3268 - accuracy: 0.9134\n",
            "Epoch 68/100\n",
            "50/50 [==============================] - 90s 2s/step - loss: 0.3371 - accuracy: 0.9109\n",
            "Epoch 69/100\n",
            "50/50 [==============================] - 89s 2s/step - loss: 0.6790 - accuracy: 0.8327\n",
            "Epoch 70/100\n",
            "50/50 [==============================] - 89s 2s/step - loss: 0.3298 - accuracy: 0.9179\n",
            "Epoch 71/100\n",
            "50/50 [==============================] - 89s 2s/step - loss: 0.3267 - accuracy: 0.9167\n",
            "Epoch 72/100\n",
            "50/50 [==============================] - 89s 2s/step - loss: 0.3393 - accuracy: 0.9125\n",
            "Epoch 73/100\n",
            "50/50 [==============================] - 93s 2s/step - loss: 0.3221 - accuracy: 0.9176\n",
            "Epoch 74/100\n",
            "50/50 [==============================] - 91s 2s/step - loss: 0.3248 - accuracy: 0.9160\n",
            "Epoch 75/100\n",
            "50/50 [==============================] - 89s 2s/step - loss: 0.3165 - accuracy: 0.9180\n",
            "Epoch 76/100\n",
            "50/50 [==============================] - 89s 2s/step - loss: 0.3202 - accuracy: 0.9181\n",
            "Epoch 77/100\n",
            "50/50 [==============================] - 91s 2s/step - loss: 0.3218 - accuracy: 0.9162\n",
            "Epoch 78/100\n",
            "50/50 [==============================] - 90s 2s/step - loss: 0.3142 - accuracy: 0.9202\n",
            "Epoch 79/100\n",
            "50/50 [==============================] - 92s 2s/step - loss: 0.3243 - accuracy: 0.9141\n",
            "Epoch 80/100\n",
            "50/50 [==============================] - 89s 2s/step - loss: 0.3167 - accuracy: 0.9198\n",
            "Epoch 81/100\n",
            "50/50 [==============================] - 91s 2s/step - loss: 0.2959 - accuracy: 0.9273\n",
            "Epoch 82/100\n",
            "50/50 [==============================] - 91s 2s/step - loss: 0.3161 - accuracy: 0.9159\n",
            "Epoch 83/100\n",
            "50/50 [==============================] - 89s 2s/step - loss: 0.2919 - accuracy: 0.9299\n",
            "Epoch 84/100\n",
            "50/50 [==============================] - 91s 2s/step - loss: 0.3100 - accuracy: 0.9183\n",
            "Epoch 85/100\n",
            "50/50 [==============================] - 89s 2s/step - loss: 0.3085 - accuracy: 0.9204\n",
            "Epoch 86/100\n",
            "50/50 [==============================] - 90s 2s/step - loss: 0.3009 - accuracy: 0.9238\n",
            "Epoch 87/100\n",
            "50/50 [==============================] - 89s 2s/step - loss: 0.2968 - accuracy: 0.9230\n",
            "Epoch 88/100\n",
            "50/50 [==============================] - 89s 2s/step - loss: 0.2991 - accuracy: 0.9226\n",
            "Epoch 89/100\n",
            "50/50 [==============================] - 89s 2s/step - loss: 0.2979 - accuracy: 0.9239\n",
            "Epoch 90/100\n",
            "50/50 [==============================] - 88s 2s/step - loss: 0.2924 - accuracy: 0.9258\n",
            "Epoch 91/100\n",
            "50/50 [==============================] - 90s 2s/step - loss: 0.3018 - accuracy: 0.9224\n",
            "Epoch 92/100\n",
            "50/50 [==============================] - 88s 2s/step - loss: 0.2829 - accuracy: 0.9311\n",
            "Epoch 93/100\n",
            "50/50 [==============================] - 89s 2s/step - loss: 0.3162 - accuracy: 0.9192\n",
            "Epoch 94/100\n",
            "50/50 [==============================] - 90s 2s/step - loss: 0.2751 - accuracy: 0.9354\n",
            "Epoch 95/100\n",
            "50/50 [==============================] - 89s 2s/step - loss: 0.2811 - accuracy: 0.9314\n",
            "Epoch 96/100\n",
            "50/50 [==============================] - 89s 2s/step - loss: 0.2846 - accuracy: 0.9292\n",
            "Epoch 97/100\n",
            "50/50 [==============================] - 89s 2s/step - loss: 0.2944 - accuracy: 0.9238\n",
            "Epoch 98/100\n",
            "50/50 [==============================] - 90s 2s/step - loss: 0.2800 - accuracy: 0.9323\n",
            "Epoch 99/100\n",
            "50/50 [==============================] - 90s 2s/step - loss: 0.2819 - accuracy: 0.9303\n",
            "Epoch 100/100\n",
            "50/50 [==============================] - 88s 2s/step - loss: 0.2766 - accuracy: 0.9320\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aV_Yoxjxfs9Y",
        "colab_type": "text"
      },
      "source": [
        "## Evaluation"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5DAScUywfwnG",
        "colab_type": "text"
      },
      "source": [
        "## Accuracy by # of moves"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7bWtbAbWce5s",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 581
        },
        "outputId": "20ae50ec-8fb9-4d6b-b635-b983263ed411"
      },
      "source": [
        "# Accuracy by number of moves\n",
        "num_batches = 3 #@param{type: \"integer\"}\n",
        "num_samples = num_batches * batch_size\n",
        "for moves in range(1, max_moves + 1):\n",
        "  print(\"For %d moves:\" % moves)\n",
        "  one_move_dataset = tf.data.Dataset.from_generator(\n",
        "    lambda: (x for x in env.generate_examples(\n",
        "      num=num_samples, possible_moves=[moves]\n",
        "    )),\n",
        "    output_types=(\n",
        "        { \"input_text\": tf.string,\n",
        "          \"start\": tf.float32,\n",
        "          \"start_name\": tf.string,\n",
        "          \"end_name\": tf.string,\n",
        "        }, tf.float32),\n",
        "    output_shapes=(\n",
        "        { \"input_text\": tf.TensorShape([]),\n",
        "          \"start\": tf.TensorShape([nqc.get_max_id('place_t')]),\n",
        "          \"start_name\": tf.TensorShape([]),\n",
        "          \"end_name\": tf.TensorShape([]),\n",
        "        }, tf.TensorShape([nqc.get_max_id('place_t')])),\n",
        "    )\n",
        "  model.evaluate(one_move_dataset.take(num_samples).batch(batch_size))\n",
        "  print(\"\")\n"
      ],
      "execution_count": 15,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "For 1 moves:\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/functional.py:543: UserWarning: Input dict contained keys ['start_name', 'end_name'] which did not match any model input. They will be ignored by the model.\n",
            "  [n for n in tensors.keys() if n not in ref_input_names])\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "3/3 [==============================] - 2s 722ms/step - loss: 0.1057 - accuracy: 1.0000\n",
            "\n",
            "For 2 moves:\n",
            "3/3 [==============================] - 2s 746ms/step - loss: 0.1099 - accuracy: 1.0000\n",
            "\n",
            "For 3 moves:\n",
            "3/3 [==============================] - 2s 731ms/step - loss: 0.1938 - accuracy: 0.9583\n",
            "\n",
            "For 4 moves:\n",
            "3/3 [==============================] - 2s 730ms/step - loss: 0.2202 - accuracy: 0.9574\n",
            "\n",
            "For 5 moves:\n",
            "3/3 [==============================] - 2s 732ms/step - loss: 0.2401 - accuracy: 0.9544\n",
            "\n",
            "For 6 moves:\n",
            "3/3 [==============================] - 2s 723ms/step - loss: 0.3043 - accuracy: 0.9287\n",
            "\n",
            "For 7 moves:\n",
            "3/3 [==============================] - 2s 749ms/step - loss: 0.3290 - accuracy: 0.9160\n",
            "\n",
            "For 8 moves:\n",
            "3/3 [==============================] - 2s 759ms/step - loss: 0.3472 - accuracy: 0.9030\n",
            "\n",
            "For 9 moves:\n",
            "3/3 [==============================] - 2s 737ms/step - loss: 0.3586 - accuracy: 0.8952\n",
            "\n",
            "For 10 moves:\n",
            "3/3 [==============================] - 2s 730ms/step - loss: 0.3465 - accuracy: 0.8975\n",
            "\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VUyK_zIZJh9w",
        "colab_type": "text"
      },
      "source": [
        "## Live Evaluation"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JFVBiLbSJqQX",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 51
        },
        "outputId": "92c2788e-ce50-47e6-fbd6-632d8d5b5686"
      },
      "source": [
        "query = \"go right then down\" #@param{type: \"string\"}\n",
        "start = \"cell_3_3\" #@param{type: \"string\"}\n",
        "\n",
        "nql_start = nqc.one(start, \"place_t\")\n",
        "if \"<UNK>\" in nql_start.eval():\n",
        "  print(\"%s is not a valid starting point. It's probably a hole or a typo.\" %\n",
        "        start)\n",
        "else:\n",
        "  example_dataset = tf.data.Dataset.from_tensors({\n",
        "      \"input_text\": query,\n",
        "      \"start\": tf.squeeze(nql_start.tf),\n",
        "  })\n",
        "  pred = model.predict(example_dataset.batch(1))\n",
        "  print(pred.shape)\n",
        "  pred_nql = nqc.as_nql(pred, \"place_t\")\n",
        "  print(pred_nql.eval(as_top=3))"
      ],
      "execution_count": 16,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "(1, 99)\n",
            "[('hole_4_4', 0.9885572), ('cell_3_5', 0.0068089347), ('cell_5_3', 0.0046146773)]\n"
          ],
          "name": "stdout"
        }
      ]
    }
  ]
}
